# coding: utf-8
# create by tongshiwei on 2018/12/9
import numpy as np
import mxnet as mx
from collections import Counter

from longling.lib.utilog import config_logging, LogLevel
from RecAgent.Agent import Agent
from GRU4Rec import *


class GRUAgent(Agent):
    def __init__(self, graph, gru4rec, soft_discount=0, path_warm_up=True, window_size=None, visit_threshold=1):
        super(GRUAgent, self).__init__(graph=graph)
        self._state = None
        self.gru4rec = gru4rec
        self.pre_path = None
        self._target = None
        self._legal_candidates = None

        self.window_size = window_size
        self.visit_threshold = visit_threshold
        self.path_warm_up = path_warm_up

        self.action_num = len(graph.nodes)
        self.soft_discount = soft_discount

    @property
    def state(self):
        return self._state[1]

    @property
    def item_value(self):
        return self._state[0]

    def begin_episode(self, exercises_record, target, *args, **kwargs):
        item_value, _state = self.gru4rec([exercise[0] for exercise in exercises_record])
        self._state = (item_value[-1], _state)
        self._target = target
        self.path = []
        self.pre_path = [e for e, c in exercises_record]
        self._legal_candidates = set(target)

    def step(self, *args, **kwargs):
        soft_discount = self.soft_discount

        action, _ = self.pi(soft_discount=soft_discount)

        return action, None

    def end_episode(self, *args, **kwargs):
        path = self.path[:]
        self._state = None
        self.pre_path = None
        self.path = None
        self._target = None
        self._legal_candidates = None

        return path

    def state_transform(self, action, correct, add_to_path=True, *args, **kwargs):
        item_value, _state = self.gru4rec([action], self.state)
        self._state = (item_value[-1], _state)
        if add_to_path is True:
            self.path.append(action)

    def get_candidate_actions(self, path=None):
        path = self.path if path is None else path
        if path:
            pnode = path[-1]
        elif self.path_warm_up is True:
            pnode = self.pre_path[-1]
        else:
            pnode = None

        visit_nodes = Counter(self.path[-self.window_size:]) if self.window_size is not None else None
        return self.graph(
            None,
            pnode=pnode,
            visit_nodes=visit_nodes,
            visit_threshold=self.visit_threshold,
            allow_shortcut=False,
            target=self._target,
            legal_candidates=self._legal_candidates,
        )

    @staticmethod
    def action_mask(a, dim):
        act = [0] * dim
        if isinstance(a, int):
            act[a] = 1
        elif isinstance(a, (list, set)):
            for _a in a:
                act[_a] = 1
        elif isinstance(a, dict):
            for k, v in a.items():
                act[k] = v
        else:
            raise TypeError

        return act

    def get_candidate_status(self, action_num, candidates, soft_candidates, soft_discount, *args, **kwargs):
        action_weight = {}

        for c in soft_candidates:
            action_weight[c] = soft_discount

        for c in candidates:
            action_weight[c] = 1

        action_value = mx.nd.array(self.item_value) * mx.nd.array(self.action_mask(action_weight, action_num))
        action_value = action_value.asnumpy().tolist()

        return list(zip(range(action_num), action_value))

    def pi(self, soft_discount=0, **kwargs):
        graph = self.graph

        # select candidates using knowledge graph
        candidates, soft_candidates = self.get_candidate_actions()

        candidate_status = self.get_candidate_status(
            action_num=self.action_num,
            candidates=candidates,
            soft_candidates=soft_candidates,
            soft_discount=soft_discount,
        )

        assert candidate_status
        candidates_set = set(candidates)
        candidate_status = [cs for cs in candidate_status if cs[0] in candidates_set]

        # value max
        action_value = max(candidate_status, key=lambda x: x[1])

        # # probability
        # actions, action_v = zip(*candidate_status)
        #
        # # softmax
        # action_v = mx.nd.array(action_v)
        # action_p = mx.nd.softmax(action_v).asnumpy()
        #
        # # linear
        # # actions, action_v = zip(*candidate_status)
        # # action_v_max, action_v_min = action_v.max(), action_v.min()
        # # action_p = (action_v - action_v_min) / (action_v_max - action_v_min)
        # # action_p = action_p / sum(action_p)
        #
        # idx = np.random.choice(range(len(action_p)), p=action_p)
        #
        # action_value = actions[idx], None

        if self.is_terminal_action(action_value[0]):
            action_value = (action_value[0], 0)

        return action_value


def main(dataset, ctx, max_steps, max_num, score_type, allow_shortcut=True, test_pre=''):
    import os
    from Graph import Graph
    from Environment import load_evaluator
    from mxnet import gpu
    from longling import wf_open

    method = 'GRU4Rec'
    if allow_shortcut is False:
        method += '-nsc'
    if score_type is True:
        method += '-score'
    root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../data/%s/" % dataset)) + os.sep
    model_dir = root_dir + "Baseline/%s/PL-%s/" % (method, max_steps)

    ctx = gpu(ctx)
    graph = Graph(dataset=dataset, disable=True)
    agent = GRUAgent(
        graph=graph,
        gru4rec=GRU4Rec(load_epoch=20, dataset=dataset),
    )

    evaluator = load_evaluator(dataset, ctx, score_type,
                               logger=config_logging(logger="evaluator", console_log_level=LogLevel.INFO))
    prefix = model_dir + "GRU4Rec%s-%s" % (test_pre, max_num)
    with wf_open(prefix + "_result.log") as wf, wf_open(prefix + "_path.log") as pwf:
        print(evaluator.eval(
            agent=agent,
            max_steps=max_steps,
            epsilon=0,
            alternative_mode=False,
            log_f=pwf,
            max_num=max_num,
            delay_estimation=False,
        ), file=wf)

    return evaluator


if __name__ == '__main__':
    import sys

    assert len(sys.argv) == 5

    dataset = sys.argv[1]
    ctx = int(sys.argv[2])
    max_steps = int(sys.argv[3])
    test_pre = sys.argv[4]

    # dataset = "irt"
    # ctx = 0
    # max_steps = 20

    max_num = 1000
    score_type = False
    allow_shortcut = False

    evaluator = main(
        dataset=dataset,
        ctx=ctx,
        max_steps=max_steps,
        max_num=max_num,
        score_type=score_type,
        allow_shortcut=allow_shortcut,
        test_pre=test_pre,
    )
